
#include <openssl/bio.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include "../iface/debug.h"
#include "../utils/BaseConvert.h"
#include "../utils/StringUtils.h"
#include "RSAPublicKey.h"

/** Constructor for a blank RSA public key. */
RSAPublicKey::RSAPublicKey()
{
   key = NULL;
   keyname = NULL;
}

/** Constructor for a blank RSA public key with a name.
 * 
 * @param name The name to attach to the blank key. */
RSAPublicKey::RSAPublicKey(const char *name)
{
   key = NULL;
   keyname = NULL;
   setKeyName(name);
}

/** Copy constructor for an RSA public key.
 * 
 * @param other The RSA public key to copy from. */
RSAPublicKey::RSAPublicKey(const RSAPublicKey &other)
{
   key = RSA_new();
   if (key == NULL)
   {
      errprint("RSA_new failed");
      return;
   }
   if (copyRSASSLPublicKey(key, other.key))
   {
      errprint("RSA_new failed");
      return;
   }
   keyname = NULL;
   setKeyName(other.keyname);
}

/** Copy constructor for an RSA public key.
 * 
 * @param other The RSA public key to copy from.
 * 
 * @return Itself. */
RSAPublicKey &RSAPublicKey::operator=(const RSAPublicKey &other)
{
   key = RSA_new();
   if (key == NULL)
   {
      errprint("RSA_new failed");
      return *this;
   }
   if (copyRSASSLPublicKey(key, other.key))
   {
      errprint("RSA_new failed");
      return *this;
   }
   keyname = NULL;
   setKeyName(other.keyname);
   return *this;
}

/** Construct an RSA public key from an RSA public key in PEM encoded
 * PKCS #8 text format.
 * 
 * @param pemtext The PEM encoded PKCS #8 text that represents an RSA
 * public key, from which this key is to be constructed.
 * 
 * @param name The name to be attached to the key, once
 * constructed. */
RSAPublicKey::RSAPublicKey(const char *pemtext, const char *name)
{
   importKeyFromASCII(pemtext, name);
}

/** Construct an RSA public key from an RSA public key in SSL format.
 * 
 * @param pubkey The SSL format RSA public key to construct this key
 * from.
 * 
 * @param name The name to be attached to the key, once
 * constructed. */
RSAPublicKey::RSAPublicKey(RSA *pubkey, const char *name)
{
   key = RSA_new();
   if (key == NULL)
   {
      errprint("RSA_new failed");
      return;
   }
   int errval = copyRSASSLPublicKey(key, pubkey);
   keyname = NULL;
   setKeyName(name);
}

/** Destructor for the RSA public key. */
RSAPublicKey::~RSAPublicKey()
{
   if (key) RSA_free(key);
   key = NULL;
   if (keyname) free((void *)keyname);
   keyname = NULL;
}

/** Get a copy of the public key as an SSL format object.
 * 
 * @param[out] pubkey A pointer to a place to put a pointer to the new
 * SSL format RSA public key which is a copy of this key. The
 * responsibility for freeing the new SSL format RSA public key is
 * transferred to the caller.
 * 
 * @return Returns zero on success, non-zero on failure. */
int RSAPublicKey::getSSLKeyComponents(RSA **pubkey) const
{
   *pubkey = RSA_new();
   if ((*pubkey) == NULL)
   {
      errprint("RSA_new failed");
      return -1;
   }
   return copyRSASSLPublicKey(*pubkey, key);
}

/** Get a copy of the public key in PEM encoded PKCS #8 text format.
 * 
 * @param[out] pemtext A pointer to a place to put a pointer to the
 * exported key in PEM encoded PKCS #8 text format. The responsibility
 * for freeing the new string is transferred to the caller.
 * 
 * @param password Password to use for unlocking the key.
 * 
 * @return Returns zero on success, non-zero on failure. */
int RSAPublicKey::getPEMKeyData(char **pemtext, const char *password) const
{
   char *data;
   int damt;
   char *dbuf = NULL;
   int retval = 0;
   BIO *bo = BIO_new(BIO_s_mem());
   if (!PEM_ASN1_write_bio((i2d_of_void *)i2d_RSAPublicKey, PEM_STRING_RSA_PUBLIC, bo, (char *)key, (password) ? EVP_des_ede3_cbc() : NULL, NULL, 0, 0, (void *)password))
   {
      retval = 1;
      errprint("could not export public key");
      SSLCheckError();
      goto bioshutdown;
   }
   SSLCheckError();
   damt = BIO_get_mem_data(bo, &data);
   *pemtext = (char *)malloc(sizeof(char)*(damt+1));
   memcpy(*pemtext, data, sizeof(char)*(damt+1));
   (*pemtext)[damt] = '\0';
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   return retval;
}

/** Export the RSA public key as the binary encoding of the SSH key
 * format in the authorized_keys file.
 * 
 * @param[out] blob A pointer to the place where the pointer to the
 * binary key blob is to be put. The responsibility for freeing the
 * new buffer is transferred to the caller.
 * 
 * @param[out] reallength A pointer to where the write the length of
 * the blob in bytes.
 * 
 * @return Returns zero on success, non-zero on failure. */
int RSAPublicKey::getSSHKeyBlob(char **blob, int *reallength) const
{
   const unsigned char keyheader[11] = { 0x00, 0x00, 0x00, 0x07, 's', 's', 'h', '-', 'r', 's', 'a' };
   unsigned char *n_buffer = NULL;
   unsigned char *e_buffer = NULL;
   int n_len = BN_num_bytes(key->n);
   int e_len = BN_num_bytes(key->e);
   n_buffer = (unsigned char *)malloc(n_len);
   if (!n_buffer) { errprint("Could not allocate buffer for pub->n."); return -1; }
   e_buffer = (unsigned char *)malloc(e_len);
   if (!e_buffer) { errprint("Could not allocate buffer for pub->e."); return -1; }
   BN_bn2bin(key->n, n_buffer);
   BN_bn2bin(key->e, e_buffer);
   int enclen = 4 + 7 + 4 + ((e_buffer[0] & 0x80) == 0x80) + e_len + 4 + ((n_buffer[0] & 0x80) == 0x80) + n_len;
   char *encode = (char *)malloc(sizeof(char)*enclen);
   memcpy(encode, keyheader, sizeof(char)*11);
   int real_len, data_start, toffset;
   toffset = 11;
   real_len = e_len + ((e_buffer[0] & 0x80) == 0x80);
   data_start = 4 + ((e_buffer[0] & 0x80) == 0x80);
   memset(&(encode[toffset]), 0, sizeof(char)*(real_len + 4));
   encode[toffset + 0] = (unsigned char)(real_len >> 24);
   encode[toffset + 1] = (unsigned char)(real_len >> 16);
   encode[toffset + 2] = (unsigned char)(real_len >>  8);
   encode[toffset + 3] = (unsigned char)(real_len);
   memcpy(&(encode[toffset + data_start]), e_buffer, sizeof(char)*e_len);
   toffset += (real_len + 4);
   real_len = n_len + ((n_buffer[0] & 0x80) == 0x80);
   data_start = 4 + ((n_buffer[0] & 0x80) == 0x80);
   memset(&(encode[toffset]), 0, sizeof(char)*(real_len + 4));
   encode[toffset + 0] = (unsigned char)(real_len >> 24);
   encode[toffset + 1] = (unsigned char)(real_len >> 16);
   encode[toffset + 2] = (unsigned char)(real_len >>  8);
   encode[toffset + 3] = (unsigned char)(real_len);
   memcpy(&(encode[toffset + data_start]), n_buffer, sizeof(char)*n_len);
   free(n_buffer); free(e_buffer);
   if (reallength) *reallength = enclen;
   if (blob) *blob = encode;
   return 0;
}

/** Export the RSA public key as a line of text that can be added to
 * the SSH authorized_keys file.
 * 
 * @param[out] textblob A pointer to the place where the pointer to
 * the text-encoded labelled key blob is to be put. This is the final
 * line that can be added to authorized_keys. The responsibility for
 * freeing the new string is transferred to the caller.
 * 
 * @param keylabel The label to go on the key on the standardised line
 * of text.
 * 
 * @return Returns zero on success, non-zero on failure. */
int RSAPublicKey::getSSHKeyBlob(char **textblob, const char *keylabel) const
{
   int enclen = 0;
   char *encode;
   if (getSSHKeyBlob(&encode, &enclen) != 0) return -1;
   char *base64string = BaseConvert::convertBytesToBase64(encode, enclen);
   free(encode);
   char *finalstring = StringUtils::varprintf("ssh-rsa %s %s", base64string, keylabel);
   free(base64string);
   if (textblob) *textblob = finalstring;
   return 0;
}

/** Compare the key to another PublicKey, and see if they are equal.
 * 
 * @param pubkey PublicKey that is to be tested to see if these are
 * equal.
 * 
 * @return 1 if they match, 0 if they do not, and negative values
 * represent errors. */
int RSAPublicKey::compareKeyEqual(PublicKey *pubkey) const
{
   if (pubkey == NULL)
   {
      errprint("comparing against NULL key");
      return 0;
   }
   switch (pubkey->getKeyType())
   {
      case KEY_TYPE_PUBLIC_RSA:
      {
         RSA *rsaKeyOther = NULL;
         ((RSAPublicKey *)pubkey)->getSSLKeyComponents(&rsaKeyOther);
         EVP_PKEY *evp1 = EVP_PKEY_new();
         EVP_PKEY *evp2 = EVP_PKEY_new();
         EVP_PKEY_set1_RSA(evp1, key);
         EVP_PKEY_set1_RSA(evp2, rsaKeyOther);
         int retval = EVP_PKEY_cmp(evp1, evp2);
         EVP_PKEY_free(evp1);
         EVP_PKEY_free(evp2);
         RSA_free(rsaKeyOther);
         SSLCheckError();
         switch (retval)
         {
            case -2:
               errprint("key compare operation not supported (EVP_PKEY_cmp())");
               break;
            case -1:
               errprint("tried to compare different key types (EVP_PKEY_cmp())");
               return 0;
            case 0:
               return 0;
            case 1:
               return 1;
         }
         return 0;
      }
      default:
      {
         return 0;
         break;
      }
   }
}

/** Copy an SSL format public key (or the public key part of a full
 * private key in SSL format, as the public key is simply a subset of
 * the private key data).
 * 
 * @param[out] dst Pointer to an initialised blank SSL format RSA key
 * to copy into.
 * 
 * @param src Pointer to an existing SSL format RSA public key to copy
 * from.
 * 
 * @return Returns zero on success, non-zero on failure. */
int RSAPublicKey::copyRSASSLPublicKey(RSA *dst, RSA *src)
{
   if (dst == NULL) { errprint("NULL public key destination"); return -1; }
   if (src == NULL) { errprint("NULL public key source"); return -1; }
   // Copy public key
   dst->n = BN_new();
   if ((src->n == NULL) || (BN_copy(dst->n, src->n) != (dst->n)))
   {
      errprint("Could not copy%s pub->n", (src->n) ? "" : " NULL");
      return -1;
   }
   dst->e = BN_new();
   if ((src->e == NULL) || (BN_copy(dst->e, src->e) != (dst->e)))
   {
      errprint("Could not copy%s pub->e", (src->e) ? "" : " NULL");
      return -1;
   }
   return 0;
}

/** Import and name an RSA public key from an RSA public key in
 * encrypted PEM format PKCS #8.
 * 
 * @param pemtext The PEM encoded PKCS #8 text that represents an RSA
 * public key, from which this key is to be constructed.
 * 
 * @param password The password to decrypt the PEM encoded PKCS #8
 * text that represents the RSA public key.
 * 
 * @return Returns zero on success, non-zero on failure. */
int RSAPublicKey::importPublicKeyWithOptionalPassword(const char *pemtext, const char *password)
{
   if (key == NULL)
      key = RSA_new();
   int retval = 0;
   BIO *bo = BIO_new_mem_buf((void *)pemtext, -1);
   if (!PEM_ASN1_read_bio((d2i_of_void *)d2i_RSAPublicKey, PEM_STRING_RSA_PUBLIC, bo, (void **)(&key), 0, (void *)password))
   {
      retval = 1;
      errprint("could not import public key");
      key = NULL;
      goto bioshutdown;
   }
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   return retval;
}

/** Import an RSA public key from some PEM format PKCS #8 text.
 * 
 * @param text The PEM encoded PKCS #8 text that represents an RSA
 * public key, from which this key is to be constructed.
 * 
 * @param name The name to be attached to the key.
 * 
 * @return Returns zero on success, non-zero on failure. */
int RSAPublicKey::importKeyFromASCII(const char *text, const char *name)
{
   setKeyName(name);
   return importPublicKeyWithOptionalPassword(text, NULL);
}

/** Import and name an RSA public key from an RSA public key in
 * encrypted PEM format PKCS #8.
 * 
 * @param text The PEM encoded PKCS #8 text that represents an RSA
 * public key, from which this key is to be constructed.
 * 
 * @param password The password to decrypt the PEM encoded PKCS #8
 * text that represents the RSA public key.
 * 
 * @param name The name to be attached to the key.
 * 
 * @return Returns zero on success, non-zero on failure. */
int RSAPublicKey::importKeyFromASCIIWithPassword(const char *text, const char *password, const char *name)
{
   setKeyName(name);
   return importPublicKeyWithOptionalPassword(text, password);
}

/** Export an RSA public key to some PEM format PKCS #8 text.
 * 
 * @param[out] text A pointer to a place to put a pointer to the
 * exported key in PEM encoded PKCS #8 text format. The responsibility
 * for freeing the new string is transferred to the caller.
 * 
 * @return Returns zero on success, non-zero on failure. */
int RSAPublicKey::exportKeyToASCII(char **text) const
{
   return getPEMKeyData(text, NULL);
}

/** Export an RSA public key to some PEM format PKCS #8 text encrypted
 * with the given password (the encryption is UNIMPLEMENTED, this is
 * plain text!).
 * 
 * @param[out] text A pointer to a place to put a pointer to the
 * exported key in encrypted PEM encoded PKCS #8 text format. The
 * responsibility for freeing the new string is transferred to the
 * caller.
 * 
 * @param password The password to encrypt and decrypt the key with.
 * 
 * @return Returns zero on success, non-zero on failure. */
int RSAPublicKey::exportKeyToASCIIWithPassword(char **text, const char *password) const
{
   return getPEMKeyData(text, password);
}
